import os 
import csv
import cv2
import pickle
import matlab.engine

import scipy.io as scio
import numpy as np

from tqdm import tqdm
from datetime import datetime
from utils.wifi_process_function import process_csi_data,csi_data_cut

"""
read label.csv get the list of labeled data
"""

file_name= "../Dataset/label.csv" 
label_list=[]
with open(file_name,newline='\n') as f:
    csv_reader=csv.reader(f)
    next(csv_reader)
    for row in csv_reader:
        label_list.append(row)


def wireless_signal_cutting(wireless_signal_raw_data,
                            wireless_signal_timestamp,
                            frame_start_timestamp,
                            frame_end_timestamp,
                            data_type="radar"):
    """
    Cuts the wireless signal data based on frame timestamps and data type.

    Args:
        wireless_signal_raw_data: The raw wireless signal data.
        wireless_signal_timestamp: Timestamps associated with the wireless signal data.
        frame_start_timestamp: The timestamp marking the start of a frame.
        frame_end_timestamp: The timestamp marking the end of a frame.
        data_type: The type of data to extract ("radar" or "wifi").

    Returns:
        clip_raw_data: The clipped raw data.
        behaviour_start_frame: The index of the start frame.
        behaviour_end_frame: The index of the end frame.
    """
    behaviour_start_frame = None
    behaviour_end_frame = None

    for ii, timestamp in enumerate(wireless_signal_timestamp):
        if round(frame_start_timestamp, 1) == round(timestamp, 1):
            behaviour_start_frame = ii
        if round(frame_end_timestamp, 1) == round(timestamp, 1):
            behaviour_end_frame = ii
            break  # 结束循环，因为已找到结束帧

    if behaviour_start_frame is None or behaviour_end_frame is None:
        raise ValueError("Frame timestamps not found in wireless_signal_timestamp")

    if data_type == "radar":
        clip_raw_data = wireless_signal_raw_data[:, :, behaviour_start_frame:behaviour_end_frame, :]
    elif data_type == "wifi":
        clip_raw_data = wireless_signal_raw_data[behaviour_start_frame:behaviour_end_frame]

    return clip_raw_data, behaviour_start_frame, behaviour_end_frame



def process_radar_data(radar_timestamp_filepath, radar_data_path,NAdc_samples,Nchirp,samplepoints,Nchannels):
    radar_timestamp = np.load(radar_timestamp_filepath, allow_pickle=True)
    radar_raw_data = scio.loadmat(radar_data_path)
    radar_raw_data = np.moveaxis(np.asarray([radar_raw_data[f'RX{i}_data'] for i in range(1, 5)]), 0, -1).reshape(NAdc_samples, Nchirp, samplepoints, Nchannels)
    return radar_timestamp, radar_raw_data



def save_data(data, save_path):
    scio.savemat(save_path, {'data': data})

def wireless_signal_clip(frame_start_timestamp,frame_end_timestamp,behaviour_name,
                         clip_radar_path,clip_csi_path,
                         radar=0,wifi=0,
                         radar_timestamp_filepath='None',radar_data_path='None',NAdc_samples=128,
                         Nchirp=64,
                         samplepoints=7500,
                         Nchannels=4,
                         csi_raw_mat_filepath='None',
                         csi_capture_start_time='None',
                         eng='None'):
    
    if radar and radar_timestamp_filepath and radar_data_path:
        radar_timestamp, radar_raw_data = process_radar_data(radar_timestamp_filepath, radar_data_path,NAdc_samples,Nchirp,samplepoints,Nchannels)
        clip_radar_raw_data, radar_behaviour_start_frame, radar_behaviour_end_frame = wireless_signal_cutting(radar_raw_data, radar_timestamp, frame_start_timestamp, frame_end_timestamp)
        radar_savepath = os.path.join(clip_radar_path, f"{radar_behaviour_start_frame}_{radar_behaviour_end_frame}_{behaviour_name}.mat")
        save_data(clip_radar_raw_data, radar_savepath)
    
    if wifi and csi_raw_mat_filepath and csi_capture_start_time:
        csi_interp_mat, csi_timestamp = process_csi_data(csi_raw_mat_filepath, csi_capture_start_time, eng)
        clip_csi_raw_data, csi_behaviour_start_frame, csi_behaviour_end_frame = wireless_signal_cutting(csi_interp_mat, csi_timestamp, frame_start_timestamp, frame_end_timestamp, data_type='wifi')
        csi_savepath = os.path.join(clip_csi_path, f"{csi_behaviour_start_frame}_{csi_behaviour_end_frame}_{behaviour_name}.mat")
        save_data(clip_csi_raw_data, csi_savepath)
    
    if not (radar or wifi):
        print("Please make sure what type of data needs to be processed!")



eng = matlab.engine.start_matlab()
vision_mmWave_csi='../Dataset/SYNCHRONIZATION/' # the root path of datasets
clip_path='../Dataset/Clip_SynC_Datasets/' # the root path of cliped datasets
vision='VIDEO'
mmWave='RADAR'
csi="WiFi"
NAdc_samples=128
Nchirp=64
samplepoints=7500
Ntxs=4
Nrxs=1


eng = matlab.engine.start_matlab()


for row in label_list:
    one_behaviour_filepath=os.path.join(vision_mmWave_csi,row[0])
    one_clip_behaviour_file_path=os.path.join(clip_path,row[0])

    if not os.path.exists(one_clip_behaviour_file_path):
        os.mkdir(one_clip_behaviour_file_path)

    # print(one_behaviour_filepath)
    frame_start_idx=int(row[1])
    frame_end_idx=int(row[2])
    behaviour_name=row[3]
    if os.path.exists(one_behaviour_filepath):
        video_path=os.path.join(one_behaviour_filepath,vision)
        radar_path=os.path.join(one_behaviour_filepath,mmWave)
        csi_path=os.path.join(one_behaviour_filepath,csi)
        
        clip_radar_path=os.path.join(one_clip_behaviour_file_path,mmWave)
        clip_csi_path=os.path.join(one_clip_behaviour_file_path,csi)

        

        if os.path.exists(video_path):

            videopath=os.path.join(video_path,
                                   os.listdir(video_path)[0],'rgb.avi')
            # video file path
            video_timestamp_filepath=os.path.join(video_path,
                                         os.listdir(video_path)[0],'frame.pkl')
            # timestamp file path 
            
            with open(video_timestamp_filepath,'rb') as f:
                timestamp=pickle.load(f)

            # print(type(timestamp[0])) <class 'datetime.datetime'>

            '''get behaviour start and end timestamps'''
            frame_start_timestamp=timestamp[frame_start_idx].timestamp()
            frame_end_timestamp=timestamp[frame_end_idx].timestamp()

            if os.path.exists(radar_path):
                if not os.path.exists(clip_radar_path):
                    os.mkdir(clip_radar_path)
                '''read radar data and timestamp'''
                radar_timestamp_filepath=os.path.join(radar_path,
                                                      os.listdir(radar_path)[0],
                                                      os.listdir(radar_path)[0]+'_frame.npy')
                radar_data_path=os.path.join(radar_path,
                                       os.listdir(radar_path)[0],
                                       os.listdir(radar_path)[0]+'.mat')
                # print(radar_data_path)
                wireless_signal_clip(frame_start_timestamp,frame_end_timestamp,behaviour_name,clip_radar_path,clip_csi_path,radar=1,
                                     radar_timestamp_filepath=radar_timestamp_filepath,radar_data_path=radar_data_path)
            if os.path.exists(csi_path):
                if not os.path.exists(clip_csi_path):
                    os.mkdir(clip_csi_path)
                '''read csi data and extract timestamp'''
                csi_capture_start_time=os.listdir(csi_path)[0]
                csi_raw_mat_filepath = os.path.join(csi_path,
                                                    csi_capture_start_time)
                wireless_signal_clip(frame_start_timestamp,frame_end_timestamp,behaviour_name,clip_radar_path,clip_csi_path,wifi=1,
                                     csi_capture_start_time=csi_capture_start_time,csi_raw_mat_filepath=csi_raw_mat_filepath,eng=eng)

    else:
        print("Error! This filepath doesn't exists")

eng.quit()